from querier.esquerier import ElasticSearchQuerier


def tokenize(b64_str, str_len=4):
    length = len(b64_str)
    n = int(length / str_len)
    out = []
    for i in range(n):
        out.append(b64_str[int(i * str_len): int((i + 1) * str_len)])
    return out


class ImgSearchQuerier(ElasticSearchQuerier):
    def __init__(self, es, index, doc_type, classifier):
        super(ImgSearchQuerier, self).__init__(es, index, doc_type)
        self.classifier = classifier

    def _build_query(self, args):
        filters = args.get('filters', {})
        if filters is None:
            filters = {}

        if not filters.get('height'):
            filters['height'] = [400]

        if not filters.get('width'):
            filters['width'] = [400]

        # order = args.get('order_by')
        from_ = args.get('from', 0)
        size_ = args.get('size', 10)

        base64_str = args.get('image_base64')
        term = args.get('term')
        if not term:
            term = ''

        classify_res = {
            "categories": [],
            "human_string": [],
            "categories_weight": []
        }

        feat2 = []

        if base64_str:
            feat, top_k, predictions, feat2, human_string = self.classifier.classify(base64_str)

            classify_res = {
                "categories": top_k.tolist(),
                "human_string": human_string,
                "categories_weight": predictions.tolist(),
            }

        filters['category'] = classify_res['categories'][0:3]
        query = self._gen_query(term, ' '.join(tokenize(feat2)), filters, from_, size_)

        return query, {}, classify_res

    def _build_result(self, es_result, param):
        total = es_result['hits']['total']

        images = []

        for hit in es_result['hits']['hits']:
            res = extract_result(hit)
            images.append(res)
        return {
            'total': total,
            'images': images,
            'categories': param['categories'],
            'human_string': param['human_string'],
            'categories_weight': param['categories_weight'],
        }

    def _gen_query(self, term, image_feature, filters, from_, size_):

        filter_clause = []
        if filters:
            self._add_filter_clause(filter_clause, filters, 'category', cond='should')
            self._add_filter_clause(filter_clause, filters, 'source', cond='should')
            self._add_filter_clause(filter_clause, filters, 'img_type', cond='should')
            self._add_filter_range_clause(filter_clause, filters, 'width')
            self._add_filter_range_clause(filter_clause, filters, 'height')
            self._add_filter_range_clause(filter_clause, filters, 'width_height_ratio')
            self._add_filter_range_clause(filter_clause, filters, 'publish_timestamp')

        query = {"query": {"bool": {}}}

        if filter_clause:
            query['query']['bool']['filter'] = filter_clause

        if image_feature.strip():
            query['query']['bool']['must'] = {
                'bool': {
                    'should': [
                        {
                            'match': {
                                'feature': {
                                    'query': image_feature,
                                    # 'boost': 20,
                                }
                            }
                        },
                    ]
                }
            }

        if term.strip():
            query['query']['bool']['must'] = {
                'bool': {
                    'should': [
                        {
                            'match': {
                                'desc': {
                                    'query': term.strip(),
                                    # 'boost': 20,
                                }
                            }
                        },
                    ]
                }
            }

        query['sort'] = []

        query['track_scores'] = True
        query['from'] = from_
        query['size'] = size_

        return query

    @staticmethod
    def _add_filter_clause(filter_clause, filters, key, cond='must'):
        if key in filters:
            if filters[key]:
                clause = []
                filter_clause.append({
                    'bool': {
                        cond: clause
                    }
                })
                for fk in filters[key]:
                    clause.append({'term': {key: fk}})
        return filter_clause

    @staticmethod
    def _add_filter_range_clause(filter_clause, filters, key):
        if key in filters:
            if filters[key]:
                clause = []
                filter_clause.append({
                    'bool': {
                        'must': clause
                    }
                })
                fk = filters[key]
                if not isinstance(fk, list) or len(fk) < 1:
                    pass
                else:
                    min_fk = fk[0]
                    if len(fk) >= 2:
                        max_fk = fk[1]
                    else:
                        max_fk = None
                    if min_fk is not None and min_fk != 'null' and min_fk != '':
                        clause.append({'range': {key: {"gte": min_fk}}})
                    if max_fk is not None and max_fk != 'null' and max_fk != '':
                        clause.append({'range': {key: {"lte": max_fk}}})
        return filter_clause


def extract_result(hit):
    source_ = hit['_source']

    source = source_['source']
    img_name = source_['img_name']
    res = {
        'image_name': img_name,
        'source': source,
        'category': source_.get('category'),
        'weibo_id': img_name.split('_')[0] if source == 'weibo' else '',
        'wechat_id': img_name.rsplit('_', 1)[0] if source == 'wechat' else '',
        'image_bae64': '',
        'image_url': get_img_url(source, source_['img_name'], source_['img_url']),
        'image_type': source_.get('img_type'),
        'score': hit['_score'],
        'desc': source_.get('desc', ''),
        'height': source_.get('height', 0),
        'width': source_.get('width', 0),
        # 'source': source_
    }
    return res


def extract_weibo_img(h):
    if h[0:4] == 'http':
        return h
    else:
        return 'http://wx3.sinaimg.cn/mw690/' + h


def get_img_url(source, img_name, min_url):
    if source == 'weibo':
        img_url = extract_weibo_img(min_url)
    elif source in ['pinterest', 'instagram']:
        img_url = 'http://localhost:5008/image/' + img_name
    elif source == 'wechat' and min_url[0:4] != 'http':
        img_url = 'http://mmbiz.qpic.cn/mmbiz_jpg/' + min_url + '/0??wx_fmt=jpeg'
    else:
        img_url = min_url

    return img_url
